home *** CD-ROM | disk | FTP | other *** search
/ Chip 2007 January, February, March & April / Chip-Cover-CD-2007-02.iso / Pakiet bezpieczenstwa / mini Pentoo LiveCD 2006.1 / mpentoo-2006.1.iso / livecd.squashfs / usr / lib / python2.4 / site-packages / Numeric / MA / MA.py < prev   
Text File  |  2006-01-20  |  73KB  |  2,138 lines
  1. """MA: a facility for dealing with missing observations
  2. MA is generally used as a Numeric.array look-alike.
  3. There are some differences in semantics, see manual.
  4. In particular note that slices are copies, not references.
  5. by Paul F. Dubois
  6.    L-264
  7.    Lawrence Livermore National Laboratory
  8.    dubois@users.sourceforge.net
  9. Copyright 1999, 2000, 2001 Regents of the University of California.
  10. Released for unlimited redistribution; see file Legal.htm
  11. Documentation is in the Numeric manual; see numpy.sourceforge.net
  12. """
  13. import Numeric
  14. import string, types, sys
  15. from Precision import *
  16. from Numeric import e, pi, NewAxis
  17. MaskType=Int0
  18. divide_tolerance = 1.e-35
  19.  
  20. class MAError (Exception):
  21.     def __init__ (self, args=None):
  22.         "Create an exception"
  23.         self.args = args
  24.     def __str__(self):
  25.         "Calculate the string representation"
  26.         return str(self.args)
  27.     __repr__ = __str__
  28.  
  29. class _MaskedPrintOption:
  30.     "One instance of this class, masked_print_option, is created."
  31.     def __init__ (self, display):
  32.         "Create the masked print option object."
  33.         self.set_display(display)
  34.         self._enabled = 1
  35.  
  36.     def display (self):
  37.         "Show what prints for masked values."
  38.         return self._display
  39.  
  40.     def set_display (self, s):
  41.         "set_display(s) sets what prints for masked values."
  42.         self._display = s
  43.  
  44.     def enabled (self):
  45.         "Is the use of the display value enabled?"
  46.         return self._enabled
  47.  
  48.     def enable(self, flag=1):
  49.         "Set the enabling flag to flag."
  50.         self._enabled = flag
  51.  
  52.     def __str__ (self):
  53.         return str(self._display)
  54.  
  55. #if you single index into a masked location you get this object.
  56. masked_print_option = _MaskedPrintOption('--')
  57.  
  58. # Use single element arrays or scalars.
  59. default_real_fill_value = Numeric.array([1.0e20]).astype(Float32)
  60. default_complex_fill_value = Numeric.array([1.0e20 + 0.0j]).astype(Complex32)
  61. default_character_fill_value = '?'
  62. default_integer_fill_value = Numeric.array([0]).astype(UnsignedInt8)
  63. default_object_fill_value = '?'
  64.  
  65. def default_fill_value (obj):
  66.     "Function to calculate default fill value for an object."
  67.     if isinstance(obj, types.FloatType):
  68.         return default_real_fill_value
  69.     elif isinstance(obj, types.IntType) or isinstance(obj, types.LongType):
  70.             return default_integer_fill_value
  71.     elif isinstance(obj, types.StringType):
  72.             return default_character_fill_value
  73.     elif isinstance(obj, types.ComplexType):
  74.             return default_complex_fill_value
  75.     elif isinstance(obj, MaskedArray) or isinstance(obj, Numeric.arraytype):
  76.         x = obj.typecode()
  77.         if x in typecodes['Float']:
  78.             return default_real_fill_value
  79.         if x in typecodes['Integer']:
  80.             return default_integer_fill_value
  81.         if x in typecodes['Complex']:
  82.             return default_complex_fill_value
  83.         if x in typecodes['Character']:
  84.             return default_character_fill_value
  85.         if x in typecodes['UnsignedInteger']:
  86.             return Numeric.absolute(default_integer_fill_value)
  87.         return default_object_fill_value
  88.     else:
  89.         return default_object_fill_value
  90.  
  91. def minimum_fill_value (obj):
  92.     "Function to calculate default fill value suitable for taking minima."
  93.     if isinstance(obj, types.FloatType):
  94.         return default_real_fill_value
  95.     elif isinstance(obj, types.IntType) or isinstance(obj, types.LongType):
  96.             return sys.maxint
  97.     elif isinstance(obj, MaskedArray) or isinstance(obj, Numeric.arraytype):
  98.         x = obj.typecode()
  99.         if x in typecodes['Float']:
  100.             return default_real_fill_value
  101.         if x in typecodes['Integer']:
  102.             return sys.maxint
  103.         if x in typecodes['UnsignedInteger']:
  104.             return sys.maxint
  105.     else:
  106.         raise TypeError, 'Unsuitable type for calculating minimum.'
  107.  
  108. def maximum_fill_value (obj):
  109.     "Function to calculate default fill value suitable for taking maxima."
  110.     if isinstance(obj, types.FloatType):
  111.         return -default_real_fill_value
  112.     elif isinstance(obj, types.IntType) or isinstance(obj, types.LongType):
  113.             return -default_integer_fill_value
  114.     elif isinstance(obj, MaskedArray) or isinstance(obj, Numeric.arraytype):
  115.         x = obj.typecode()
  116.         if x in typecodes['Float']:
  117.             return -default_real_fill_value
  118.         if x in typecodes['Integer']:
  119.             return -sys.maxint
  120.         if x in typecodes['UnsignedInteger']:
  121.             return 0
  122.     else:
  123.         raise TypeError, 'Unsuitable type for calculating maximum.'
  124.  
  125. def set_fill_value (a, fill_value):
  126.     "Set fill value of a if it is a masked array."
  127.     if isMaskedArray(a): 
  128.         a.set_fill_value (fill_value)
  129.  
  130. def getmask (a):
  131.     """Mask of values in a; could be None.
  132.        Returns None if a is not a masked array.
  133.        To get an array for sure use getmaskarray."""
  134.     if isinstance(a, MaskedArray): 
  135.         return a.raw_mask()
  136.     else: 
  137.         return None
  138.  
  139. def getmaskarray (a):
  140.     """Mask of values in a; an array of zeros if mask is None
  141.      or not a masked array. Caution: has savespace attribute,
  142.      and is a byte-sized integer.
  143.      Do not try to add up entries, for example.
  144.     """
  145.     m = getmask(a)
  146.     if m is None:
  147.         return make_mask_none(shape(a))
  148.     else:
  149.         return m
  150.  
  151. def is_mask (m):
  152.     """Is m a legal mask? Does not check contents, only type.
  153.     """
  154.     if m is None or (isinstance(m, Numeric.ArrayType) and \
  155.                      m.typecode() == MaskType):
  156.         return 1
  157.     else:
  158.         return 0
  159.    
  160. def make_mask (m, copy=0, flag=0):
  161.     """make_mask(m, copy=0, flag=0)
  162.        return m as a mask, creating a copy if necessary or requested.
  163.        Can accept any sequence of integers or None. Does not check 
  164.        that contents must be 0s and 1s.
  165.        if flag, return None if m contains no true elements.
  166.     """
  167.     if m is None: 
  168.         return None
  169.     elif isinstance(m, Numeric.ArrayType):
  170.         if m.typecode() == MaskType:
  171.             if copy:
  172.                 result = Numeric.array(m, savespace=1)
  173.             else:
  174.                 result = Numeric.array(m, copy=0, savespace=1)
  175.         else:
  176.             result = m.astype(MaskType)
  177.             result.savespace(1)
  178.     else:
  179.         result = Numeric.array(filled(m,1), MaskType, savespace=1)
  180.  
  181.     if flag and not Numeric.sometrue(Numeric.ravel(result)):
  182.         return None
  183.     else:
  184.         return result
  185.  
  186. def make_mask_none (s):
  187.     "Return a mask of all zeros of shape s."
  188.     result = Numeric.zeros(s, MaskType)
  189.     result.savespace(1)
  190.     result.shape = s
  191.     return result
  192.  
  193. create_mask = make_mask_none #backwards compatibility
  194.  
  195. def mask_or (m1, m2):
  196.     """Logical or of the mask candidates m1 and m2, treating None as false.
  197.        Result may equal m1 or m2 if the other is None.
  198.      """
  199.     if m1 is None: return make_mask(m2)
  200.     if m2 is None: return make_mask(m1)
  201.     if m1 is m2 and is_mask(m1): return m1
  202.     return make_mask(Numeric.logical_or(m1, m2))
  203.  
  204. def filled (a, value = None):
  205.     """a as a contiguous Numeric array with any masked areas replaced by value
  206.     if value is None or the special element "masked", fill_value(a)
  207.     is used instead. 
  208.  
  209.     If a is already a contiguous Numeric array, a itself is returned.
  210.  
  211.     filled(a) can be used to be sure that the result is Numeric when 
  212.     passing an object a to other software ignorant of MA, in particular to
  213.     Numeric itself.
  214.     """
  215.     if isinstance(a, MaskedArray):
  216.         return a.filled(value)
  217.     elif isinstance(a, Numeric.ArrayType) and a.iscontiguous():
  218.         return a
  219.     elif isinstance(a, types.DictType):
  220.         return Numeric.array(a, 'O')
  221.     else:
  222.         return Numeric.array(a)
  223.  
  224. def fill_value (a):
  225.     """
  226.     The fill value of a, if it has one; otherwise, the default fill value
  227.     for that type.
  228.     """
  229.     if isMaskedArray(a):
  230.         result = a.fill_value()
  231.     else:
  232.         result = default_fill_value(a)
  233.     return result
  234.  
  235. def common_fill_value (a, b):
  236.     "The common fill_value of a and b, if there is one, or None"
  237.     t1 = fill_value(a)
  238.     t2 = fill_value(b)
  239.     if t1 == t2: return t1
  240.     return None
  241.  
  242. # Domain functions return 1 where the argument(s) are not in the domain.
  243. class domain_check_interval:
  244.     "domain_check_interval(a,b)(x) = true where x < a or y > b"
  245.     def __init__(self, y1, y2):
  246.         "domain_check_interval(a,b)(x) = true where x < a or y > b"
  247.         self.y1 = y1
  248.         self.y2 = y2
  249.  
  250.     def __call__ (self, x):
  251.         "Execute the call behavior."
  252.         return Numeric.logical_or(Numeric.greater (x, self.y2),
  253.                                    Numeric.less(x, self.y1)
  254.                                   )
  255.  
  256. class domain_tan:
  257.     "domain_tan(eps) = true where abs(cos(x)) < eps)"
  258.     def __init__(self, eps):
  259.         "domain_tan(eps) = true where abs(cos(x)) < eps)"
  260.         self.eps = eps
  261.  
  262.     def __call__ (self, x):
  263.         "Execute the call behavior."
  264.         return Numeric.less(Numeric.absolute(Numeric.cos(x)), self.eps)
  265.  
  266. class domain_greater:
  267.     "domain_greater(v)(x) = true where x <= v"
  268.     def __init__(self, critical_value):
  269.         "domain_greater(v)(x) = true where x <= v"
  270.         self.critical_value = critical_value
  271.  
  272.     def __call__ (self, x):
  273.         "Execute the call behavior."
  274.         return Numeric.less_equal (x, self.critical_value)
  275.  
  276. class domain_greater_equal:
  277.     "domain_greater_equal(v)(x) = true where x < v"
  278.     def __init__(self, critical_value):
  279.         "domain_greater_equal(v)(x) = true where x < v"
  280.         self.critical_value = critical_value
  281.  
  282.     def __call__ (self, x):
  283.         "Execute the call behavior."
  284.         return Numeric.less (x, self.critical_value)
  285.  
  286. class masked_unary_operation:
  287.     def __init__ (self, aufunc, fill=0, domain=None):
  288.         """ masked_unary_operation(aufunc, fill=0, domain=None)
  289.             aufunc(fill) must be defined
  290.             self(x) returns aufunc(x) 
  291.             with masked values where domain(x) is true or getmask(x) is true. 
  292.         """
  293.         self.f = aufunc
  294.         self.fill = fill
  295.         self.domain = domain
  296.         self.__doc__ = getattr(aufunc, "__doc__", str(aufunc))
  297.  
  298.     def __call__ (self, a, *args, **kwargs):
  299.         "Execute the call behavior."
  300. # Numeric tries to return scalars rather than arrays when given scalars.
  301.         m = getmask(a)
  302.         d1 = filled(a, self.fill)
  303.         if self.domain is not None:
  304.             m = mask_or(m, self.domain(d1))
  305.         if m is None:
  306.             result = self.f(d1, *args, **kwargs)
  307.             if type(result) is Numeric.ArrayType:
  308.                 return masked_array (result)
  309.             else:
  310.                 return result
  311.         else:
  312.             dx = masked_array(d1, m)
  313.             result = self.f(filled(dx, self.fill), *args, **kwargs)
  314.             if type(result) is Numeric.ArrayType:
  315.                 return masked_array(result, m)
  316.             elif m[...]:
  317.                 return masked
  318.             else:
  319.                 return result
  320.  
  321.     def __str__ (self):
  322.         return "Masked version of " + str(self.f)
  323.  
  324.  
  325. class domain_safe_divide:
  326.     def __init__ (self, tolerance=divide_tolerance):
  327.         self.tolerance = tolerance
  328.     def __call__ (self, a, b):
  329.         return Numeric.absolute(a) * self.tolerance >= Numeric.absolute(b)
  330.  
  331. class domained_binary_operation:
  332.     """Binary operations that have a domain, like divide. These are complicated so they
  333.        are a separate class. They have no reduce, outer or accumulate.
  334.     """
  335.     def __init__ (self, abfunc, domain, fillx=0, filly=0):
  336.         """abfunc(fillx, filly) must be defined.
  337.            abfunc(x, filly) = x for all x to enable reduce.
  338.         """
  339.         self.f = abfunc
  340.         self.domain = domain
  341.         self.fillx = fillx
  342.         self.filly = filly
  343.         self.__doc__ = getattr(abfunc, "__doc__", str(abfunc))
  344.  
  345.     def __call__(self, a, b):
  346.         "Execute the call behavior."
  347.         ma = getmask(a)
  348.         mb = getmask(b)
  349.         d1 = filled(a, self.fillx)
  350.         d2 = filled(b, self.filly)
  351.         t = self.domain(d1, d2)
  352.         
  353.         if Numeric.sometrue(t):
  354.             d2 = where(t, self.filly, d2)
  355.             mb = mask_or(mb, t)
  356.         m = mask_or(ma, mb)
  357.         if m is None:
  358.             result =  self.f(d1, d2)
  359.             if type(result) is Numeric.ArrayType:
  360.                 return masked_array(result)
  361.             else:
  362.                 return result
  363.         result = self.f(d1, d2)
  364.         if type(result) is Numeric.ArrayType:
  365.             if m.shape != result.shape:
  366.                 m = mask_or(getmaskarray(a), getmaskarray(b))
  367.             return masked_array(result, m)
  368.         elif m[...]:
  369.             return masked
  370.         else:
  371.             return result
  372.     def __str__ (self):
  373.         return "Masked version of " + str(self.f)
  374.  
  375. class masked_binary_operation:
  376.     def __init__ (self, abfunc, fillx=0, filly=0):
  377.         """abfunc(fillx, filly) must be defined.
  378.            abfunc(x, filly) = x for all x to enable reduce.
  379.         """
  380.         self.f = abfunc
  381.         self.fillx = fillx
  382.         self.filly = filly
  383.         self.__doc__ = getattr(abfunc, "__doc__", str(abfunc))
  384.  
  385.     def __call__ (self, a, b, *args, **kwargs):
  386.         "Execute the call behavior."
  387.         m = mask_or(getmask(a), getmask(b))
  388.         if m is None:
  389.             d1 = filled(a, self.fillx)
  390.             d2 = filled(b, self.filly)
  391.             result =  self.f(d1, d2, *args, **kwargs)
  392.             if type(result) is Numeric.ArrayType:
  393.                 return masked_array(result)
  394.             else:
  395.                 return result
  396.         d1 = filled(a, self.fillx)
  397.         d2 = filled(b, self.filly)
  398.         result = self.f(d1, d2, *args, **kwargs)
  399.         if type(result) is Numeric.ArrayType:
  400.             if m.shape != result.shape:
  401.                 m = mask_or(getmaskarray(a), getmaskarray(b))
  402.             return masked_array(result, m)
  403.         elif m[...]:
  404.             return masked
  405.         else:
  406.             return result
  407.  
  408.     def reduce (self, target, axis=0):
  409.         """Reduce target along the given axis with this function."""
  410.         m = getmask(target)
  411.         t = filled(target, self.filly)
  412.         if t.shape == ():
  413.             t.shape = (1,)
  414.             if m is not None:
  415.                m = make_mask(m, copy=1)
  416.                m.shape = (1,)
  417.         if m is None:
  418.             return masked_array (self.f.reduce (t, axis))
  419.         else:
  420.             t = masked_array (t, m)
  421.             t = self.f.reduce(filled(t, self.filly), axis)
  422.             m = Numeric.logical_and.reduce(m, axis)
  423.             if isinstance(t, Numeric.ArrayType):
  424.                 return masked_array(t, m, fill_value(target))
  425.             elif m:
  426.                 return masked
  427.             else:
  428.                 return t
  429.  
  430.     def outer (self, a, b):
  431.         "Return the function applied to the outer product of a and b."
  432.         ma = getmask(a)
  433.         mb = getmask(b)
  434.         if ma is None and mb is None:
  435.             m = None
  436.         else:
  437.             ma = getmaskarray(a)
  438.             mb = getmaskarray(b)
  439.             m = logical_or.outer(ma, mb)
  440.         d = self.f.outer(filled(a, self.fillx), filled(b, self.filly))
  441.         return masked_array(d, m)
  442.       
  443.     def accumulate (self, target, axis=0):
  444.         """Accumulate target along axis after filling with y fill value."""
  445.         t = filled(target, self.filly)
  446.         return masked_array (self.f.accumulate (t, axis))
  447.     def __str__ (self):
  448.         return "Masked version of " + str(self.f)
  449.  
  450. sqrt = masked_unary_operation(Numeric.sqrt, 0.0, domain_greater_equal(0.0))
  451. log = masked_unary_operation(Numeric.log, 1.0, domain_greater(0.0))
  452. log10 = masked_unary_operation(Numeric.log10, 1.0, domain_greater(0.0))
  453. exp = masked_unary_operation(Numeric.exp)
  454. conjugate = masked_unary_operation(Numeric.conjugate)
  455. sin = masked_unary_operation(Numeric.sin)
  456. cos = masked_unary_operation(Numeric.cos)
  457. tan = masked_unary_operation(Numeric.tan, 0.0, domain_tan(1.e-35))
  458. arcsin = masked_unary_operation(Numeric.arcsin, 0.0, domain_check_interval(-1.0, 1.0))
  459. arccos = masked_unary_operation(Numeric.arccos, 0.0, domain_check_interval(-1.0, 1.0))
  460. arctan = masked_unary_operation(Numeric.arctan)
  461. # Missing from Numeric
  462. # arcsinh = masked_unary_operation(Numeric.arcsinh)
  463. # arccosh = masked_unary_operation(Numeric.arccosh)
  464. # arctanh = masked_unary_operation(Numeric.arctanh)
  465. sinh = masked_unary_operation(Numeric.sinh)
  466. cosh = masked_unary_operation(Numeric.cosh)
  467. tanh = masked_unary_operation(Numeric.tanh)
  468. absolute = masked_unary_operation(Numeric.absolute)
  469. fabs = masked_unary_operation(Numeric.fabs)
  470. negative = masked_unary_operation(Numeric.negative)
  471. nonzero = masked_unary_operation(Numeric.nonzero)
  472. around = masked_unary_operation(Numeric.around)
  473. floor = masked_unary_operation(Numeric.floor)
  474. ceil = masked_unary_operation(Numeric.ceil)
  475. sometrue = masked_unary_operation(Numeric.sometrue)
  476. alltrue = masked_unary_operation(Numeric.alltrue, 1)
  477. logical_not = masked_unary_operation(Numeric.logical_not)
  478.  
  479. add = masked_binary_operation(Numeric.add)
  480. subtract = masked_binary_operation(Numeric.subtract)
  481. subtract.reduce = None
  482. multiply = masked_binary_operation(Numeric.multiply, 1, 1)
  483. divide = domained_binary_operation(Numeric.divide, domain_safe_divide(), 0, 1)
  484. true_divide = domained_binary_operation(Numeric.true_divide, domain_safe_divide(), 0, 1)
  485. floor_divide = domained_binary_operation(Numeric.floor_divide, domain_safe_divide(), 0, 1)
  486. remainder = domained_binary_operation(Numeric.remainder, domain_safe_divide(), 0, 1)
  487. fmod = domained_binary_operation(Numeric.fmod, domain_safe_divide(), 0, 1)
  488. hypot = masked_binary_operation(Numeric.hypot)
  489. arctan2 = masked_binary_operation(Numeric.arctan2, 0.0, 1.0)
  490. arctan2.reduce = None
  491. equal = masked_binary_operation(Numeric.equal)
  492. equal.reduce = None
  493. not_equal = masked_binary_operation(Numeric.not_equal)
  494. not_equal.reduce = None
  495. less_equal = masked_binary_operation(Numeric.less_equal)
  496. less_equal.reduce = None
  497. greater_equal = masked_binary_operation(Numeric.greater_equal)
  498. greater_equal.reduce = None
  499. less = masked_binary_operation(Numeric.less)
  500. less.reduce = None
  501. greater = masked_binary_operation(Numeric.greater)
  502. greater.reduce = None
  503. logical_and = masked_binary_operation(Numeric.logical_and)
  504. logical_or = masked_binary_operation(Numeric.logical_or)
  505. logical_xor = masked_binary_operation(Numeric.logical_xor)
  506. bitwise_and = masked_binary_operation(Numeric.bitwise_and)
  507. bitwise_or = masked_binary_operation(Numeric.bitwise_or)
  508. bitwise_xor = masked_binary_operation(Numeric.bitwise_xor)
  509. rank = Numeric.rank
  510. shape = Numeric.shape
  511. size = Numeric.size
  512.  
  513. class MaskedArray (object):
  514.     """Arrays with possibly masked values. 
  515.        Masked values of 1 exclude element from the computation.
  516.        Construction:
  517.            x = array(data, typecode=None, copy=1, savespace=0,
  518.                      mask = None, fill_value=None)
  519.  
  520.        If copy=0, every effort is made not to copy the data:
  521.            If data is a MaskedArray, and argument mask=None,
  522.            then the candidate data is data.raw_data() and the
  523.            mask used is data.mask(). If data is a Numeric array,
  524.            it is used as the candidate raw data.
  525.            If savespace != data.spacesaver() or typecode is not None and
  526.            is != data.typecode() then a data copy is required.  
  527.            Otherwise, the candidate is used.
  528.  
  529.        If a data copy is required, raw data stored is the result of:
  530.        Numeric.array(data, typecode=typecode, copy=copy, savespace=savespace)
  531.  
  532.        If mask is None there are no masked values. Otherwise mask must
  533.        be convertible to an array of integers of typecode MaskType, 
  534.        with values 1 or 0, and of the same shape as x.
  535.  
  536.        fill_value is used to fill in masked values when necessary, 
  537.        such as when printing and in method/function filled().
  538.        The fill_value is not used for computation within this module.
  539.  
  540.        If savespace is 1, the data is given the spacesaver property, and
  541.        the mask is replaced by None if all its elements are true.
  542.     """
  543.     handler_cache_key = 'MA.MaskedArray'
  544.  
  545.     def __init__(self, data, typecode=None, 
  546.                   copy=1, savespace=None, mask=None, fill_value=None,
  547.                   ):
  548.         """array(data, typecode=None,copy=1, savespace=None,
  549.                     mask=None, fill_value=None)
  550.            If data already a Numeric array, its typecode and spacesaver()
  551.            become the default values for typecode and savespace.
  552.         """
  553.         tc = typecode
  554.         ss = savespace
  555.         need_data_copied = copy
  556.         if isinstance(data, MaskedArray): 
  557.             c = data.raw_data()
  558.             ctc = c.typecode()
  559.             if tc is None:
  560.                 tc = ctc
  561.             elif tc != ctc:
  562.                 need_data_copied = 1
  563.             css = c.spacesaver()
  564.             if ss is None:
  565.                 ss = css
  566.             elif ss != css:
  567.                 need_data_copied = 1
  568.             else:
  569.                 ss = 0
  570.             if mask is None:
  571.                 mask = data.mask()
  572.             elif mask is not None: #attempting to change the mask
  573.                 need_data_copied = 1
  574.  
  575.         elif isinstance(data, Numeric.ArrayType): 
  576.             c = data
  577.             ctc = c.typecode()
  578.             if tc is None:
  579.                 tc = ctc
  580.             elif tc != ctc:
  581.                 need_data_copied = 1
  582.             css = c.spacesaver()
  583.             if ss is None:
  584.                 ss = css
  585.             elif ss != css:
  586.                 need_data_copied = 1
  587.         else:
  588.             need_data_copied = 0 #because I'll do it now
  589.             if ss is None: 
  590.                 ss = 0
  591.             c = Numeric.array(data, tc, savespace=ss)
  592.  
  593.         if need_data_copied:
  594.             if tc == ctc:
  595.                 self._data = Numeric.array(c, copy=1, savespace = ss)
  596.             else:
  597.                 self._data = c.astype(tc)
  598.                 self._data.savespace(ss)
  599.         else:
  600.             self._data = c
  601.  
  602.         if mask is None:
  603.             self._mask = None
  604.             self._shared_mask = 0
  605.         else:
  606.             self._mask = make_mask (mask, flag=ss)
  607.             if self._mask is None:
  608.                 self._shared_mask = 0
  609.             else:
  610.                 self._shared_mask = (self._mask is mask)
  611.                 nm = size(self._mask)
  612.                 nd = size(self._data)
  613.                 if nm != nd:
  614.                     if nm == 1:
  615.                         self._mask = Numeric.resize(self._mask, self._data.shape)
  616.                         self._shared_mask = 0
  617.                     elif nd == 1:
  618.                         self._data = Numeric.resize(self._data, self._mask.shape)
  619.                         self._data.shape = self._mask.shape
  620.                     else:
  621.                         raise MAError, "Mask and data not compatible."
  622.                 elif nm == 1 and shape(self._mask) != shape(self._data):
  623.                     self.unshare_mask()
  624.                     self._mask.shape = self._data.shape
  625.  
  626.         self.set_fill_value(fill_value)
  627.  
  628.     def __array__ (self, t = None):
  629.         "Special hook for Numeric. Converts to Numeric if possible."
  630.         if self._mask is not None:
  631.             if Numeric.sometrue(Numeric.ravel(self._mask)):
  632.                 raise MAError, \
  633.                 """Cannot automatically convert masked array to Numeric because data 
  634.                    is masked in one or more locations.
  635.                 """
  636.             else:  # Mask is all false
  637.                    # Optimize to avoid future invocations of this section.
  638.                 self._mask = None
  639.                 self._shared_mask = 0
  640.         if t: 
  641.             return self._data.astype(t)
  642.         else:
  643.             return self._data
  644.     
  645.     def _get_shape(self):
  646.         "Return the current shape."
  647.         return self._data.shape
  648.  
  649.     def _set_shape (self, newshape):
  650.         "Set the array's shape."
  651.         if not self._data.iscontiguous():
  652.             self._data = Numeric.array(self._data, self._data.typecode(), 
  653.                                         1, self._data.spacesaver())
  654.         self._data.shape = newshape
  655.         if self._mask is not None:
  656.             self.unshare_mask()
  657.             if not self._mask.iscontiguous():
  658.                 self._mask = Numeric.array(self._mask, MaskType, 1, 1)
  659.             self._mask.shape = newshape
  660.             
  661.     def _get_flat(self):
  662.         """Calculate the flat value. 
  663.         """
  664.         if self._mask is None:
  665.             return masked_array(self._data.flat, mask=None,
  666.                                 fill_value = self.fill_value())
  667.         else:
  668.             return masked_array(self._data.flat, 
  669.                                 mask=self._mask.flat,
  670.                                 fill_value = self.fill_value())
  671.  
  672.     def _set_flat (self, value):
  673.         "x.flat = value"
  674.         y = self.flat
  675.         y[:] = value
  676.  
  677.     def _get_real(self):
  678.         "Get the real part of a complex array."
  679.         if self._mask is None:
  680.             return masked_array(self._data.real, mask=None,
  681.                             fill_value = self.fill_value())
  682.         else:
  683.             return masked_array(self._data.real, mask=self.mask().flat,
  684.                             fill_value = self.fill_value())
  685.     
  686.     def _set_real (self, value):
  687.         "x.real = value"
  688.         y = self.real
  689.         y[...] = value
  690.  
  691.     def _get_imaginary(self):
  692.         "Get the imaginary part of a complex array."
  693.         if self._mask is None:
  694.             return masked_array(self._data.imaginary, mask=None,
  695.                             fill_value = self.fill_value())
  696.         else:
  697.             return masked_array(self._data.imaginary, mask=self.mask().flat,
  698.                             fill_value = self.fill_value())
  699.  
  700.     def _set_imaginary (self, value):
  701.         "x.imaginary = value"
  702.         y = self.imaginary
  703.         y[...] = value
  704.  
  705.     def __str__(self):
  706.         """Calculate the str representation, using masked for fill if 
  707.            it is enabled. Otherwise fill with fill value.
  708.         """
  709.         if masked_print_option.enabled():
  710.             f = masked_print_option
  711.         else:
  712.             f = self.fill_value()
  713.         return str(filled(self, f))
  714.  
  715.     def __repr__(self):
  716.         """Calculate the repr representation, using masked for fill if 
  717.            it is enabled. Otherwise fill with fill value.
  718.         """
  719.         with_mask = """\
  720. array(data = 
  721.  %(data)s,
  722.       mask = 
  723.  %(mask)s,
  724.       fill_value=%(fill)s)
  725. """
  726.         with_mask1 = """\
  727. array(data = %(data)s,
  728.       mask = %(mask)s,
  729.       fill_value=%(fill)s)
  730. """
  731.         without_mask = """array(
  732.  %(data)s)""" 
  733.         without_mask1 = """array(%(data)s)""" 
  734.  
  735.         n = len(self.shape)
  736.         if self._mask is None:
  737.             if n <=1:
  738.                 return without_mask1 % {'data':str(self.filled())}
  739.             return without_mask % {'data':str(self.filled())}
  740.         else:
  741.             if n <=1:
  742.                 return with_mask % {
  743.                     'data': str(self.filled()),
  744.                     'mask': str(self.mask()),
  745.                     'fill': str(self.fill_value())
  746.                     }
  747.             return with_mask % {
  748.                 'data': str(self.filled()),
  749.                 'mask': str(self.mask()),
  750.                 'fill': str(self.fill_value())
  751.                 }
  752.         without_mask1 = """array(%(data)s)""" 
  753.         if self._mask is None:
  754.             return without_mask % {'data':str(self.filled())}
  755.         else:
  756.             return with_mask % {
  757.                 'data': str(self.filled()),
  758.                 'mask': str(self.mask()),
  759.                 'fill': str(self.fill_value())
  760.                 }
  761.  
  762.     def __float__(self):
  763.         "Convert self to float."
  764.         self.unmask()
  765.         if self._mask is not None:
  766.             raise MAError, 'Cannot convert masked element to a Python float.'
  767.         return float(self.raw_data()[...])
  768.  
  769.     def __int__(self):
  770.         "Convert self to int."
  771.         self.unmask()
  772.         if self._mask is not None:
  773.             raise MAError, 'Cannot convert masked element to a Python int.'
  774.         return int(self.raw_data()[...])
  775.  
  776. # Note copy semantics here differ from Numeric        
  777.     def __getitem__(self, i): 
  778.         "Get copy of item described by i."
  779.         m = self._mask
  780.         dout = self._data[i]
  781.         ss = self._data.spacesaver()
  782.         tc =self._data.typecode()
  783.         if type(dout) is Numeric.ArrayType:
  784.             if m is None:
  785.                 result = array(dout, typecode=tc, copy = 1, savespace=ss)
  786.             else:
  787.                 result = array(dout, typecode=tc, copy = 1, savespace=ss, 
  788.                           mask = m[i], fill_value=self.fill_value())
  789.             return result
  790.         elif m is None or not m[i]:
  791.             return dout  #scalar
  792.         else:  #scalar but masked
  793.             return masked
  794.  
  795.     def __getslice__(self, i, j): 
  796.         "Get copy of slice described by i, j"
  797.         m = self._mask
  798.         dout = self._data[i:j]
  799.         ss = self._data.spacesaver()
  800.         tc =self._data.typecode()
  801.         if m is None:
  802.             return array(dout, typecode=tc, copy = 1, savespace=ss)
  803.         else:
  804.             return array(dout, typecode=tc, copy = 1, savespace=ss, 
  805.                       mask = m[i:j], fill_value=self.fill_value())
  806.             
  807. # --------
  808. # setitem and setslice notes
  809. # note that if value is masked, it means to mask those locations.
  810. # setting a value changes the mask to match the value in those locations.
  811.     def __setitem__(self, index, value):
  812.         "Set item described by index. If value is masked, mask those locations."
  813.         if self is masked: 
  814.             raise MAError, 'Cannot alter the masked element.'
  815.         if value is masked:
  816.             if self._mask is None:
  817.                 self._mask = make_mask_none(self._data.shape)
  818.                 self._shared_mask = 0
  819.             else:
  820.                 self.unshare_mask()
  821.             self._mask[index] = 1
  822.             return
  823.         m = getmask(value)
  824.         value = filled(value).astype(self._data.typecode())
  825.         self._data[index] = value
  826.         if m is None:
  827.             if self._mask is not None:
  828.                 self.unshare_mask()
  829.                 self._mask[index] = 0
  830.         else:
  831.             if self._mask is None:
  832.                 self._mask = make_mask_none(self._data.shape)
  833.                 self._shared_mask = 0
  834.             else:
  835.                 self.unshare_mask()
  836.             self._mask[index] = m
  837.  
  838.     def __setslice__(self, i, j, value):
  839.         "Set slice i:j; if value is masked, mask those locations."
  840.         if self is masked: 
  841.             raise MAError, 'Cannot alter the masked element.'
  842.         if value is masked:
  843.             if self._mask is None:
  844.                 self._mask = make_mask_none(self._data.shape)
  845.                 self._shared_mask = 0
  846.             self._mask[i:j] = 1
  847.             return
  848.         m = getmask(value)
  849.         value = filled(value).astype(self._data.typecode())
  850.         self._data[i:j] = value
  851.         if m is None:
  852.             if self._mask is not None:
  853.                 self.unshare_mask()
  854.                 self._mask[i:j] = 0
  855.         else:
  856.             if self._mask is None:
  857.                 self._mask = make_mask_none(self._data.shape)
  858.                 self._shared_mask = 0
  859.             self._mask[i:j] = m
  860.  
  861.     def __len__ (self):
  862.         """Return length of first dimension. This is weird but Python's
  863.          slicing behavior depends on it."""
  864.         return len(self._data)
  865.  
  866.     def __and__(self, other):
  867.         "Return bitwise_and"
  868.         return bitwise_and(self, other)
  869.  
  870.     def __or__(self, other):
  871.         "Return bitwise_or"
  872.         return bitwise_or(self, other)
  873.  
  874.     def __xor__(self, other):
  875.         "Return bitwise_xor"
  876.         return bitwise_xor(self, other)
  877.  
  878.     __rand__ = __and__
  879.     __ror__ = __or__
  880.     __rxor__ = __xor__
  881.  
  882.     def __abs__(self): 
  883.         "Return absolute(self)"
  884.         return absolute(self)
  885.  
  886.     def __neg__(self): 
  887.         "Return negative(self)"
  888.         return negative(self)
  889.  
  890.     def __pos__(self):
  891.         "Return array(self)"
  892.         return array(self)
  893.  
  894.     def __add__(self, other): 
  895.         "Return add(self, other)"
  896.         return add(self, other)
  897.  
  898.     __radd__ = __add__
  899.  
  900.     def __mod__ (self, other):
  901.         "Return remainder(self, other)"
  902.         return remainder(self, other)
  903.  
  904.     def __rmod__ (self, other):
  905.         "Return remainder(other, self)"
  906.         return remainder(other, self)
  907.  
  908.     def __lshift__ (self, n):
  909.         return left_shift(self, n)
  910.  
  911.     def __rshift__ (self, n):
  912.         return right_shift(self, n)
  913.                         
  914.     def __sub__(self, other): 
  915.         "Return subtract(self, other)"
  916.         return subtract(self, other)
  917.  
  918.     def __rsub__(self, other): 
  919.         "Return subtract(other, self)"
  920.         return subtract(other, self)
  921.  
  922.     def __mul__(self, other): 
  923.         "Return multiply(self, other)"
  924.         return multiply(self, other)
  925.     
  926.     __rmul__ = __mul__
  927.  
  928.     def __div__(self, other): 
  929.         "Return divide(self, other)"
  930.         return divide(self, other)
  931.  
  932.     def __rdiv__(self, other): 
  933.         "Return divide(other, self)"
  934.         return divide(other, self)
  935.  
  936.     def __truediv__(self, other): 
  937.         "Return divide(self, other)"
  938.         return true_divide(self, other)
  939.  
  940.     def __rtruediv__(self, other): 
  941.         "Return divide(other, self)"
  942.         return true_divide(other, self)
  943.  
  944.     def __floordiv__(self, other): 
  945.         "Return divide(self, other)"
  946.         return floor_divide(self, other)
  947.  
  948.     def __rfloordiv__(self, other): 
  949.         "Return divide(other, self)"
  950.         return floor_divide(other, self)
  951.  
  952.     def __pow__(self,other, third=None): 
  953.         "Return power(self, other, third)"
  954.         return power(self, other, third)
  955.  
  956.     def __sqrt__(self): 
  957.         "Return sqrt(self)"
  958.         return sqrt(self)
  959.  
  960.     def __iadd__(self, other): 
  961.         "Add other to self in place."
  962.         t = self._data.typecode()
  963.         f = filled(other,0)
  964.         t1 = f.typecode()
  965.         if t == t1:
  966.             pass
  967.         elif t in typecodes['Integer']:
  968.             if t1 in typecodes['Integer']:
  969.                 f = f.astype(t)
  970.             else:
  971.                 raise TypeError, 'Incorrect type for in-place operation.'
  972.         elif t in typecodes['Float']:
  973.             if t1 in typecodes['Integer']:
  974.                 f = f.astype(t)
  975.             elif t1 in typecodes['Float']:
  976.                 f = f.astype(t)
  977.             else:
  978.                 raise TypeError, 'Incorrect type for in-place operation.'
  979.         elif t in typecodes['Complex']:
  980.             if t1 in typecodes['Integer']:
  981.                 f = f.astype(t)
  982.             elif t1 in typecodes['Float']:
  983.                 f = f.astype(t)
  984.             elif t1 in typecodes['Complex']:
  985.                 f = f.astype(t)
  986.             else:
  987.                 raise TypeError, 'Incorrect type for in-place operation.'
  988.         else:
  989.             raise TypeError, 'Incorrect type for in-place operation.'
  990.  
  991.         if self._mask is None:
  992.             self._data += f
  993.             m = getmask(other)
  994.             self._mask = m
  995.             self._shared_mask = m is not None
  996.         else:
  997.             result = add(self, masked_array(f, mask=getmask(other)))
  998.             self._data = result.raw_data()
  999.             self._mask = result.raw_mask()
  1000.             self._shared_mask = 1
  1001.         return self
  1002.  
  1003.     def __imul__(self, other): 
  1004.         "Add other to self in place."
  1005.         t = self._data.typecode()
  1006.         f = filled(other,0)
  1007.         t1 = f.typecode()
  1008.         if t == t1:
  1009.             pass
  1010.         elif t in typecodes['Integer']:
  1011.             if t1 in typecodes['Integer']:
  1012.                 f = f.astype(t)
  1013.             else:
  1014.                 raise TypeError, 'Incorrect type for in-place operation.'
  1015.         elif t in typecodes['Float']:
  1016.             if t1 in typecodes['Integer']:
  1017.                 f = f.astype(t)
  1018.             elif t1 in typecodes['Float']:
  1019.                 f = f.astype(t)
  1020.             else:
  1021.                 raise TypeError, 'Incorrect type for in-place operation.'
  1022.         elif t in typecodes['Complex']:
  1023.             if t1 in typecodes['Integer']:
  1024.                 f = f.astype(t)
  1025.             elif t1 in typecodes['Float']:
  1026.                 f = f.astype(t)
  1027.             elif t1 in typecodes['Complex']:
  1028.                 f = f.astype(t)
  1029.             else:
  1030.                 raise TypeError, 'Incorrect type for in-place operation.'
  1031.         else:
  1032.             raise TypeError, 'Incorrect type for in-place operation.'
  1033.  
  1034.         if self._mask is None:
  1035.             self._data *= f
  1036.             m = getmask(other)
  1037.             self._mask = m
  1038.             self._shared_mask = m is not None
  1039.         else:
  1040.             result = multiply(self, masked_array(f, mask=getmask(other)))
  1041.             self._data = result.raw_data()
  1042.             self._mask = result.raw_mask()
  1043.             self._shared_mask = 1
  1044.         return self
  1045.  
  1046.     def __isub__(self, other): 
  1047.         "Subtract other from self in place."
  1048.         t = self._data.typecode()
  1049.         f = filled(other,0)
  1050.         t1 = f.typecode()
  1051.         if t == t1:
  1052.             pass
  1053.         elif t in typecodes['Integer']:
  1054.             if t1 in typecodes['Integer']:
  1055.                 f = f.astype(t)
  1056.             else:
  1057.                 raise TypeError, 'Incorrect type for in-place operation.'
  1058.         elif t in typecodes['Float']:
  1059.             if t1 in typecodes['Integer']:
  1060.                 f = f.astype(t)
  1061.             elif t1 in typecodes['Float']:
  1062.                 f = f.astype(t)
  1063.             else:
  1064.                 raise TypeError, 'Incorrect type for in-place operation.'
  1065.         elif t in typecodes['Complex']:
  1066.             if t1 in typecodes['Integer']:
  1067.                 f = f.astype(t)
  1068.             elif t1 in typecodes['Float']:
  1069.                 f = f.astype(t)
  1070.             elif t1 in typecodes['Complex']:
  1071.                 f = f.astype(t)
  1072.             else:
  1073.                 raise TypeError, 'Incorrect type for in-place operation.'
  1074.         else:
  1075.             raise TypeError, 'Incorrect type for in-place operation.'
  1076.  
  1077.         if self._mask is None:
  1078.             self._data -= f
  1079.             m = getmask(other)
  1080.             self._mask = m
  1081.             self._shared_mask = m is not None
  1082.         else:
  1083.             result = subtract(self, masked_array(f, mask=getmask(other)))
  1084.             self._data = result.raw_data()
  1085.             self._mask = result.raw_mask()
  1086.             self._shared_mask = 1
  1087.         return self
  1088.  
  1089.  
  1090.  
  1091.     def __idiv__(self, other): 
  1092.         "Divide self by other in place."
  1093.         t = self._data.typecode()
  1094.         f = filled(other,0)
  1095.         t1 = f.typecode()
  1096.         if t == t1:
  1097.             pass
  1098.         elif t in typecodes['Integer']:
  1099.             if t1 in typecodes['Integer']:
  1100.                 f = f.astype(t)
  1101.             else:
  1102.                 raise TypeError, 'Incorrect type for in-place operation.'
  1103.         elif t in typecodes['Float']:
  1104.             if t1 in typecodes['Integer']:
  1105.                 f = f.astype(t)
  1106.             elif t1 in typecodes['Float']:
  1107.                 f = f.astype(t)
  1108.             else:
  1109.                 raise TypeError, 'Incorrect type for in-place operation.'
  1110.         elif t in typecodes['Complex']:
  1111.             if t1 in typecodes['Integer']:
  1112.                 f = f.astype(t)
  1113.             elif t1 in typecodes['Float']:
  1114.                 f = f.astype(t)
  1115.             elif t1 in typecodes['Complex']:
  1116.                 f = f.astype(t)
  1117.             else:
  1118.                 raise TypeError, 'Incorrect type for in-place operation.'
  1119.         else:
  1120.             raise TypeError, 'Incorrect type for in-place operation.'
  1121.         mo = getmask(other)
  1122.         result = divide(self, masked_array(f, mask=mo))
  1123.         self._data = result.raw_data()
  1124.         dm = result.raw_mask()
  1125.         if dm is not self._mask:
  1126.             self._mask = dm
  1127.             self._shared_mask = 1
  1128.         return self
  1129.  
  1130.     def __eq__(self,other): 
  1131.         return equal(self,other)
  1132.  
  1133.     def __ne__(self,other): 
  1134.         return not_equal(self,other)
  1135.  
  1136.     def __lt__(self,other): 
  1137.         return less(self,other)
  1138.  
  1139.     def __le__(self,other): 
  1140.         return less_equal(self,other)
  1141.  
  1142.     def __gt__(self,other): 
  1143.         return greater(self,other)
  1144.  
  1145.     def __ge__(self,other): 
  1146.         return greater_equal(self,other)
  1147.  
  1148.     def astype (self, tc):
  1149.         "return self as array of given type."
  1150.         d = self._data.astype(tc)
  1151.         d.savespace(self._data.spacesaver())
  1152.         return array(d, mask=self._mask)
  1153.  
  1154.     def byte_swapped(self):
  1155.         """Returns the raw data field, byte_swapped. Included for consistency
  1156.          with Numeric but doesn't make sense in this context.
  1157.         """
  1158.         return self._data.byte_swapped()
  1159.  
  1160.     def compressed (self):
  1161.         "A 1-D array of all the non-masked data."
  1162.         d = Numeric.ravel(self._data)
  1163.         if self._mask is None:
  1164.             return array(d)
  1165.         else:
  1166.             m = 1 - Numeric.ravel(self._mask)
  1167.             c = Numeric.compress(m, d)
  1168.             return array(c, copy=0)
  1169.  
  1170.     def count (self, axis = None):
  1171.         "Count of the non-masked elements in a, or along a certain axis."
  1172.         m = self._mask
  1173.         s = self._data.shape
  1174.         ls = len(s)
  1175.         if m is None:
  1176.             if ls == 0:
  1177.                 return 1
  1178.             if ls == 1:
  1179.                 return s[0]
  1180.             if axis is None:
  1181.                 return reduce(lambda x,y:x*y, s)
  1182.             else:
  1183.                 n = s[axis]
  1184.                 t = list(s)
  1185.                 del t[axis]
  1186.                 return ones(t) * n
  1187.         if axis is None:
  1188.             w = Numeric.ravel(m).astype(Int)  #avoid savespace truncation
  1189.             n1 = size(w)
  1190.             if n1 == 1:
  1191.                  n2 = w[0]
  1192.             else:
  1193.                  n2 = Numeric.add.reduce(w)
  1194.             return n1 - n2
  1195.         else:
  1196.             n1 = size(m, axis)
  1197.             n2 = sum(m.astype(Int), axis)
  1198.             return n1 - n2
  1199.  
  1200.     def dot (self, other):
  1201.         "s.dot(other) = innerproduct(s, other)"
  1202.         return innerproduct(self, other)
  1203.  
  1204.     def fill_value(self):
  1205.         "Get the current fill value."
  1206.         return self._fill_value
  1207.  
  1208.     def filled (self, fill_value=None):
  1209.         """A Numeric array with masked values filled. If fill_value is None,
  1210.            use self.fill_value(). 
  1211.            
  1212.            If mask is None, copy data only if not contiguous.
  1213.            Result is always a contiguous, Numeric array.
  1214.         """
  1215.         d = self._data
  1216.         m = self._mask
  1217.         if m is None:
  1218.             if d.iscontiguous():
  1219.                 return d
  1220.             else:
  1221.                 return Numeric.array(d, typecode=d.typecode(), copy=1,
  1222.                                        savespace = d.spacesaver())
  1223.         value = fill_value
  1224.         if value is None:
  1225.             value = self._fill_value
  1226.         if self is masked:
  1227.             result = Numeric.array(value)
  1228.             result.shape = d.shape
  1229.         else:
  1230.             try:
  1231.                 result = Numeric.array(d, typecode=d.typecode(), copy=1,
  1232.                                            savespace = d.spacesaver())
  1233.                 Numeric.putmask(result, m, value)
  1234.             except:
  1235.                 result = Numeric.choose(m, (d, value))
  1236.         return result
  1237.  
  1238.     def ids (self):
  1239.         """Return the ids of the data and mask areas"""
  1240.         return (id(self._data), id(self._mask))
  1241.  
  1242.     def iscontiguous (self):
  1243.         "Is the data contiguous?"
  1244.         return self._data.iscontiguous()
  1245.  
  1246.     def itemsize(self):
  1247.         "Item size of each data item."
  1248.         return self._data.itemsize()
  1249.  
  1250.     def mask(self):
  1251.         "Return the data mask, or None. Result contiguous."
  1252.         m = self._mask
  1253.         if m is None: 
  1254.             return m
  1255.         elif m.iscontiguous():
  1256.             return m
  1257.         else:
  1258.             return Numeric.array(self._mask)
  1259.     
  1260.     def outer(self, other):
  1261.         "s.outer(other) = outerproduct(s, other)"
  1262.         return outerproduct(self, other)
  1263.  
  1264.     def put (self, values):
  1265.         """Set the non-masked entries of self to filled(values). 
  1266.            No change to mask
  1267.         """
  1268.         iota = Numeric.arange(self.size())
  1269.         if self._mask is None:
  1270.             ind = iota
  1271.         else:
  1272.             ind = Numeric.compress(1 - self._mask, iota)
  1273.         if len(ind) != size(values):
  1274.             raise MAError, "x.put(values) incorrect count of values."       
  1275.         Numeric.put (self._data, ind, filled(values))
  1276.  
  1277.     def putmask (self, values):
  1278.         """Set the masked entries of self to filled(values). 
  1279.            Mask changed to None.
  1280.         """
  1281.         if self._mask is not None:
  1282.             iota = Numeric.arange(self.size())
  1283.             ind = Numeric.compress(self._mask, iota)
  1284.             if len(ind) != size(values):
  1285.                 raise MAError, "x.put(values) incorrect count of values." 
  1286.             Numeric.put (self._data, ind, filled(values))
  1287.             self._mask = None
  1288.             self._shared_mask = 0
  1289.  
  1290.     def raw_data (self):
  1291.         """ The raw data; portions may be meaningless. 
  1292.             May be noncontiguous. Expert use only."""
  1293.         return self._data
  1294.  
  1295.     def raw_mask (self):
  1296.         """ The raw mask; portions may be meaningless. 
  1297.             May be noncontiguous. Expert use only.
  1298.         """
  1299.         return self._mask
  1300.  
  1301.     def spacesaver (self):
  1302.         "Get the spacesaver attribute."
  1303.         return self._data.spacesaver()
  1304.         
  1305.     def savespace (self, value):
  1306.         "Set the spacesaver attribute to value"
  1307.         self._data.savespace(value)
  1308.  
  1309.     def set_fill_value (self, v=None):
  1310.         "Set the fill value to v. Omit v to restore default."
  1311.         if v is None:
  1312.             v = default_fill_value (self.raw_data())
  1313.         self._fill_value = v
  1314.  
  1315.     def size (self, axis = None):
  1316.         "Number of elements in array, or in a particular axis."
  1317.         s = self._data.shape
  1318.         if axis is None:
  1319.             if len(s) == 0:
  1320.                 return 1
  1321.             else:
  1322.                 return reduce(lambda x,y: x*y, s)
  1323.         else:
  1324.             return s[axis]
  1325.         
  1326.     def spacesaver (self):
  1327.         "spacesaver() queries the spacesaver attribute."
  1328.         return self._data.spacesaver()
  1329.  
  1330.     def typecode(self):
  1331.         return self._data.typecode()
  1332.  
  1333.     def tolist(self, fill_value=None):
  1334.         "Convert to list"
  1335.         return self.filled(fill_value).tolist()
  1336.  
  1337.     def tostring(self, fill_value=None):
  1338.         "Convert to string"
  1339.         return self.filled(fill_value).tostring()
  1340.  
  1341.     def unmask (self):
  1342.         "Replace the mask by None if possible."
  1343.         if self._mask is None: return
  1344.         m = make_mask(self._mask, flag=1)
  1345.         if m is None:
  1346.             self._mask = None
  1347.             self._shared_mask = 0
  1348.  
  1349.     def unshare_mask (self):
  1350.         "If currently sharing mask, make a copy."
  1351.         if self._shared_mask:
  1352.             self._mask = make_mask (self._mask, copy=1, flag=0)
  1353.             self._shared_mask = 0
  1354.  
  1355.     shape = property(_get_shape, _set_shape,
  1356.            doc = 'tuple giving the shape of the array')
  1357.  
  1358.     flat = property(_get_flat, _set_flat,
  1359.            doc = 'Access array in flat form.')
  1360.  
  1361.     real = property(_get_real, _set_real,
  1362.            doc = 'Access the real part of the array')
  1363.  
  1364.     imaginary = property(_get_imaginary, _set_imaginary,
  1365.            doc = 'Access the imaginary part of the array')
  1366.  
  1367.     imag = imaginary
  1368.                                                         
  1369. array = MaskedArray
  1370.  
  1371. class MaskedScalar (MaskedArray):
  1372.     def __init__ (self):
  1373.         MaskedArray.__init__ (self, [0], mask=[1])
  1374.         self._data.shape = ()
  1375.         self._mask.shape = ()
  1376.  
  1377.     shape = property(MaskedArray._get_shape)
  1378.  
  1379. masked = MaskedScalar()
  1380.  
  1381. def isMaskedArray (x):
  1382.     "Is x a masked array, that is, an instance of MaskedArray?"
  1383.     return isinstance(x, MaskedArray)
  1384.  
  1385. isarray = isMaskedArray
  1386. isMA = isMaskedArray  #backward compatibility
  1387.  
  1388. def allclose (a, b, fill_value=1, rtol=1.e-5, atol=1.e-8):
  1389.     """ Returns true if all components of a and b are equal
  1390.         subject to given tolerances.
  1391.         If fill_value is 1, masked values considered equal.
  1392.         If fill_value is 0, masked values considered unequal.
  1393.         The relative error rtol should be positive and << 1.0
  1394.         The absolute error atol comes into play for those elements
  1395.         of b that are very small or zero; it says how small a must be also.
  1396.     """
  1397.     m = mask_or(getmask(a), getmask(b))
  1398.     d1 = filled(a)
  1399.     d2 = filled(b)
  1400.     x = filled(array(d1, copy=0, mask=m), fill_value).astype(Float)
  1401.     y = filled(array(d2, copy=0, mask=m), 1).astype(Float)
  1402.     d = Numeric.less_equal(Numeric.absolute(x-y), atol + rtol * Numeric.absolute(y))
  1403.     return Numeric.alltrue(Numeric.ravel(d))    
  1404.  
  1405. def allequal (a, b, fill_value=1):
  1406.     """
  1407.         True if all entries of  a and b are equal, using 
  1408.         fill_value as a truth value where either or both are masked.
  1409.     """
  1410.     m = mask_or(getmask(a), getmask(b))
  1411.     if m is None:
  1412.         x = filled(a)
  1413.         y = filled(b)
  1414.         d = Numeric.equal(x, y)
  1415.         return Numeric.alltrue(Numeric.ravel(d))
  1416.     elif fill_value:
  1417.         x = filled(a)
  1418.         y = filled(b)
  1419.         d = Numeric.equal(x, y)
  1420.         dm = array(d, mask=m, copy=0)
  1421.         return Numeric.alltrue(Numeric.ravel(filled(dm, 1)))
  1422.     else:
  1423.         return 0
  1424.     
  1425. def masked_values (data, value, rtol=1.e-5, atol=1.e-8, copy=1,
  1426.     savespace=0): 
  1427.     """
  1428.        masked_values(data, value, rtol=1.e-5, atol=1.e-8)
  1429.        Create a masked array; mask is None if possible.
  1430.        If copy==0, and otherwise possible, result
  1431.        may share data values with original array.
  1432.        Let d = filled(data, value). Returns d 
  1433.        masked where abs(data-value)<= atol + rtol * abs(value)
  1434.        if d is of a floating point type. Otherwise returns
  1435.        masked_object(d, value, copy, savespace)
  1436.     """
  1437.     abs = Numeric.absolute
  1438.     d = filled(data, value)
  1439.     if d.typecode() in typecodes['Float']:
  1440.         m = Numeric.less_equal(abs(d-value), atol+rtol*abs(value))    
  1441.         m = make_mask(m, flag=1)
  1442.         return array(d, mask = m, savespace=savespace, copy=copy, 
  1443.                       fill_value=value)
  1444.     else:
  1445.         return masked_object(d, value, copy, savespace)
  1446.     
  1447. def masked_object (data, value, copy=1, savespace=0):
  1448.     "Create array masked where exactly data equal to value"
  1449.     d = filled(data, value)
  1450.     dm = make_mask(Numeric.equal(d, value), flag=1)
  1451.     return array(d, mask=dm, copy=copy, savespace=savespace, 
  1452.                    fill_value=value)
  1453.     
  1454. def arrayrange(start, stop=None, step=1, typecode=None):
  1455.     """Just like range() except it returns a array whose type can be specfied
  1456.     by the keyword argument typecode.
  1457.     """
  1458.     return array(Numeric.arrayrange(start, stop, step, typecode))
  1459.  
  1460. arange = arrayrange   
  1461.  
  1462. def fromstring (s, t):
  1463.     "Construct a masked array from a string. Result will have no mask."
  1464.     return masked_array(Numeric.fromstring(s, t))
  1465.  
  1466. def left_shift (a, n):
  1467.     "Left shift n bits"
  1468.     m = getmask(a)
  1469.     if m is None:
  1470.         d = Numeric.left_shift(filled(a), n)
  1471.         return masked_array(d)
  1472.     else:
  1473.         d = Numeric.left_shift(filled(a,0), n)
  1474.         return masked_array(d, m)
  1475.  
  1476. def right_shift (a, n):
  1477.     "Right shift n bits"
  1478.     m = getmask(a)
  1479.     if m is None:
  1480.         d = Numeric.right_shift(filled(a), n)
  1481.         return masked_array(d)
  1482.     else:
  1483.         d = Numeric.right_shift(filled(a,0), n)
  1484.         return masked_array(d, m)
  1485.  
  1486. def resize (a, new_shape):
  1487.     """resize(a, new_shape) returns a new array with the specified shape.
  1488.     The original array's total size can be any size."""
  1489.     m = getmask(a)
  1490.     if m is not None:
  1491.         m = Numeric.resize(m, new_shape)
  1492.     result = array(Numeric.resize(filled(a), new_shape), mask=m)
  1493.     result.set_fill_value(fill_value(a))
  1494.     return result
  1495.  
  1496. def repeat(a, repeats, axis=0):
  1497.     """repeat elements of a repeats times along axis
  1498.        repeats is a sequence of length a.shape[axis]
  1499.        telling how many times to repeat each element.
  1500.     """
  1501.     af = filled(a)
  1502.     if isinstance(repeats, types.IntType):
  1503.         repeats = tuple([repeats]*(shape(af)[axis]))
  1504.     
  1505.     m = getmask(a)
  1506.     if m is not None:
  1507.         m = Numeric.repeat(m, repeats, axis)
  1508.     d = Numeric.repeat(af, repeats, axis)
  1509.     result = masked_array(d, m)
  1510.     result.set_fill_value(fill_value(a))
  1511.     return result
  1512.  
  1513. def identity(n):
  1514.     """identity(n) returns the identity matrix of shape n x n.
  1515.     """
  1516.     return array(Numeric.identity(n))
  1517.  
  1518. def indices (dimensions, typecode=None):
  1519.     """indices(dimensions,typecode=None) returns an array representing a grid
  1520.     of indices with row-only, and column-only variation.
  1521.     """
  1522.     return array(Numeric.indices(dimensions, typecode))
  1523.  
  1524. def zeros (shape, typecode=Int, savespace=0):
  1525.     """zeros(n, typecode=Int, savespace=0) = 
  1526.      an array of all zeros of the given length or shape."""
  1527.     return array(Numeric.zeros(shape, typecode, savespace))
  1528.     
  1529. def ones (shape, typecode=Int, savespace=0):
  1530.     """ones(n, typecode=Int, savespace=0) = 
  1531.      an array of all ones of the given length or shape."""
  1532.     return array(Numeric.ones(shape, typecode, savespace))
  1533.     
  1534.  
  1535. def count (a, axis = None):
  1536.     "Count of the non-masked elements in a, or along a certain axis."   
  1537.     a = masked_array(a)
  1538.     return a.count(axis)
  1539.  
  1540. def power (a, b, third=None):
  1541.     "a**b"
  1542.     if third is not None:
  1543.         raise MAError, "3-argument power not supported."
  1544.     ma = getmask(a)
  1545.     mb = getmask(b)
  1546.     m = mask_or(ma, mb)
  1547.     fa = filled(a, 1)
  1548.     fb = filled(b, 1)
  1549.     if fb.typecode() in typecodes["Integer"]:
  1550.         return masked_array(Numeric.power(fa, fb), m)
  1551.     md = make_mask(Numeric.less_equal (fa, 0), flag=1)
  1552.     m = mask_or(m, md)
  1553.     if m is None:
  1554.         return masked_array(Numeric.power(fa, fb))
  1555.     else:
  1556.         fa = Numeric.where(m, 1, fa)
  1557.         return masked_array(Numeric.power(fa, fb), m)
  1558.        
  1559. def masked_array (a, mask=None, fill_value=None):
  1560.     """masked_array(a, mask=None) = 
  1561.        array(a, mask=mask, copy=0, fill_value=fill_value)
  1562.        Use fill_value(a) if None.
  1563.     """
  1564. #
  1565. # This is an unfortunate copy of what is in fill_value
  1566. # but I want the name fill_value as a parameter.
  1567. # 
  1568.     if fill_value is None: 
  1569.         if isMaskedArray(a):
  1570.             fill_value = a.fill_value()
  1571.         else:
  1572.             fill_value = default_fill_value(a)
  1573.     return array(a, mask=mask, copy=0, fill_value=fill_value) 
  1574.  
  1575. sum = add.reduce
  1576. product = multiply.reduce
  1577.  
  1578. def average (a, axis=0, weights=None, returned = 0):
  1579.     """average(a, axis=0, weights=None)
  1580.        Computes average along indicated axis. 
  1581.        If axis is None, average over the entire array
  1582.        Inputs can be integer or floating types; result is of type Float.
  1583.        
  1584.        If weights are given, result is sum(a*weights)/(sum(weights)*1.0)
  1585.        weights must have a's shape or be the 1-d with length the size
  1586.        of a in the given axis.
  1587.  
  1588.        If returned, return a tuple: the result and the sum of the weights 
  1589.        or count of values. Results will have the same shape.
  1590.  
  1591.        masked values in the weights will be set to 0.0
  1592.     """
  1593.     a = masked_array(a)
  1594.     mask = a.mask()
  1595.     ash = a.shape
  1596.     if ash == ():
  1597.         ash = (1,)
  1598.     if axis is None:
  1599.         if mask is None:
  1600.             if weights is None:
  1601.                 n = add.reduce(a.raw_data().flat)
  1602.                 d = reduce(lambda x, y: x * y, ash, 1.0)
  1603.             else:
  1604.                 w = filled(weights, 0.0).flat
  1605.                 n = Numeric.add.reduce(a.raw_data().flat * w)
  1606.                 d = Numeric.add.reduce(w)
  1607.                 del w
  1608.         else:
  1609.             if weights is None:
  1610.                 n = add.reduce(a.flat)
  1611.                 w = Numeric.choose(mask, (1.0,0.0)).flat
  1612.                 d = Numeric.add.reduce(w)
  1613.                 del w
  1614.             else:
  1615.                 w = array(filled(weights, 0.0), Numeric.Float, mask=mask).flat
  1616.                 n = add.reduce(a.flat * w)
  1617.                 d = add.reduce(w)
  1618.                 del w
  1619.     else:
  1620.         if mask is None:
  1621.             if weights is None:
  1622.                 d = ash[axis] * 1.0
  1623.                 n = Numeric.add.reduce(a.raw_data(), axis)
  1624.             else:
  1625.                 w = filled(weights, 0.0) 
  1626.                 wsh = w.shape
  1627.                 if wsh == ():
  1628.                     wsh = (1,)
  1629.                 if wsh == ash:
  1630.                     w = Numeric.array(w, Float, copy=0)
  1631.                     n = add.reduce(a*w, axis)
  1632.                     d = add.reduce(w, axis) 
  1633.                     del w
  1634.                 elif wsh == (ash[axis],):
  1635.                     ni = ash[axis]
  1636.                     r = [NewAxis]*len(ash)
  1637.                     r[axis] = slice(None,None,1)
  1638.                     w = eval ("w["+ repr(tuple(r)) + "] * ones(ash, Float)")
  1639.                     n = add.reduce(a*w, axis)
  1640.                     d = add.reduce(w, axis)
  1641.                     del w, r
  1642.                 else:
  1643.                     raise ValueError, 'average: weights wrong shape.'
  1644.         else:
  1645.             if weights is None:
  1646.                 n = add.reduce(a, axis)
  1647.                 w = Numeric.choose(mask, (1.0, 0.0))
  1648.                 d = Numeric.add.reduce(w, axis)
  1649.                 del w
  1650.             else:
  1651.                 w = filled(weights, 0.0) 
  1652.                 wsh = w.shape
  1653.                 if wsh == ():
  1654.                     wsh = (1,)
  1655.                 if wsh == ash:
  1656.                     w = array(w, Float, mask=mask, copy=0)
  1657.                     n = add.reduce(a*w, axis)
  1658.                     d = add.reduce(w, axis) 
  1659.                 elif wsh == (ash[axis],):
  1660.                     ni = ash[axis]
  1661.                     r = [NewAxis]*len(ash)
  1662.                     r[axis] = slice(None,None,1)
  1663.                     w = eval ("w["+ repr(tuple(r)) + "] * masked_array(ones(ash, Float), mask)")
  1664.                     n = add.reduce(a*w, axis)
  1665.                     d = add.reduce(w, axis)
  1666.                 else:
  1667.                     raise ValueError, 'average: weights wrong shape.'
  1668.                 del w
  1669.     # print n, d, repr(mask), repr(weights)
  1670.     result = divide (n, d)
  1671.     del n
  1672.  
  1673.     if isinstance(result, MaskedArray):
  1674.         result.unmask() 
  1675.         if returned:  
  1676.             if not isinstance(d, MaskedArray):
  1677.                 d = masked_array(d)
  1678.             if not d.shape == result.shape:
  1679.                 d = ones(result.shape, Float) * d
  1680.             d.unmask()
  1681.     if returned:
  1682.         return result, d
  1683.     else:
  1684.         return result
  1685.  
  1686. def where (condition, x, y):
  1687.     """where(condition, x, y) is x where condition is nonzero, y otherwise.
  1688.        condition must be convertible to an integer array.
  1689.        Answer is always the shape of condition.
  1690.        The type depends on x and y. It is integer if both x and y are 
  1691.        the value masked.
  1692.     """ 
  1693.     fc = filled(not_equal(condition,0), 0)
  1694.     if x is masked:
  1695.         xv = 0
  1696.         xm = 1
  1697.     else:
  1698.         xv = filled(x)
  1699.         xm = getmask(x)
  1700.         if xm is None: xm = 0
  1701.     if y is masked:
  1702.         yv = 0
  1703.         ym = 1
  1704.     else:
  1705.         yv = filled(y)
  1706.         ym = getmask(y)
  1707.         if ym is None: ym = 0
  1708.     d = Numeric.choose(fc, (yv, xv))
  1709.     md = Numeric.choose(fc, (ym, xm))
  1710.     m = getmask(condition)
  1711.     m = make_mask(mask_or(m, md), copy=0, flag=1)
  1712.     return masked_array(d, m)
  1713.  
  1714. def choose (indices, t):
  1715.     "Returns array shaped like indices with elements chosen from t"
  1716.     def fmask (x):
  1717.         if x is masked: return 1
  1718.         return filled(x)
  1719.     def nmask (x):
  1720.         if x is masked: return 1
  1721.         m = getmask(x)
  1722.         if m is None: return 0
  1723.         return m
  1724.     c = filled(indices,0)
  1725.     masks = [nmask(x) for x in t]
  1726.     a = [fmask(x) for x in t]
  1727.     d = Numeric.choose(c, a)
  1728.     m = Numeric.choose(c, masks)
  1729.     m = make_mask(mask_or(m, getmask(indices)), copy=0, flag=1)
  1730.     return masked_array(d, m)
  1731.  
  1732. def masked_where(condition, x, copy=1):
  1733.     """Return x as an array masked where condition is true. 
  1734.        Also masked where x or condition masked.
  1735.     """
  1736.     cm = filled(condition,1)
  1737.     m = mask_or(getmask(x), cm)
  1738.     return array(filled(x), copy=copy, mask=m)
  1739.  
  1740. def masked_greater(x, value, copy=1):
  1741.     "masked_greater(x, value) = x masked where x > value"
  1742.     return masked_where(greater(x, value), x, copy)
  1743.  
  1744. def masked_greater_equal(x, value, copy=1):
  1745.     "masked_greater_equal(x, value) = x masked where x >= value"
  1746.     return masked_where(greater_equal(x, value), x, copy)
  1747.  
  1748. def masked_less(x, value, copy=1):
  1749.     "masked_less(x, value) = x masked where x < value"
  1750.     return masked_where(less(x, value), x, copy)
  1751.  
  1752. def masked_less_equal(x, value, copy=1):
  1753.     "masked_less_equal(x, value) = x masked where x <= value"
  1754.     return masked_where(less_equal(x, value), x, copy)
  1755.  
  1756. def masked_not_equal(x, value, copy=1):
  1757.     "masked_not_equal(x, value) = x masked where x != value"
  1758.     d = filled(x,0)
  1759.     c = Numeric.not_equal(d, value)
  1760.     m = mask_or(c, getmask(x))
  1761.     return array(d, mask=m, copy=copy)
  1762.  
  1763. def masked_equal(x, value, copy=1):
  1764.     """masked_equal(x, value) = x masked where x == value
  1765.        For floating point consider masked_values(x, value) instead.
  1766.     """
  1767.     d = filled(x,0)
  1768.     c = Numeric.equal(d, value)
  1769.     m = mask_or(c, getmask(x))
  1770.     return array(d, mask=m, copy=copy)
  1771.  
  1772. def masked_inside(x, v1, v2, copy=1):
  1773.     """x with mask of all values of x that are inside [v1,v2]
  1774.        v1 and v2 can be given in either order.
  1775.     """
  1776.     if v2 < v1:
  1777.         t = v2
  1778.         v2 = v1
  1779.         v1 = t
  1780.     d=filled(x, 0)
  1781.     c = Numeric.logical_and(Numeric.less_equal(d, v2), Numeric.greater_equal(d, v1))
  1782.     m = mask_or(c, getmask(x))
  1783.     return array(d, mask = m, copy=copy)
  1784.  
  1785. def masked_outside(x, v1, v2, copy=1):
  1786.     """x with mask of all values of x that are outside [v1,v2]
  1787.        v1 and v2 can be given in either order.
  1788.     """
  1789.     if v2 < v1:
  1790.         t = v2
  1791.         v2 = v1
  1792.         v1 = t
  1793.     d = filled(x,0)
  1794.     c = Numeric.logical_or(Numeric.less(d, v1), Numeric.greater(d, v2))
  1795.     m = mask_or(c, getmask(x))
  1796.     return array(d, mask = m, copy=copy)
  1797.  
  1798. def reshape (a, newshape):
  1799.     "Copy of a with a new shape."
  1800.     m = getmask(a)
  1801.     d = Numeric.reshape(filled(a), newshape)
  1802.     if m is None:
  1803.         return masked_array(d)
  1804.     else:
  1805.         return masked_array(d, mask=Numeric.reshape(m, newshape))
  1806.  
  1807. def ravel (a):
  1808.     "a as one-dimensional, may share data and mask"
  1809.     m = getmask(a)
  1810.     d = Numeric.ravel(filled(a))   
  1811.     if m is None:
  1812.         return masked_array(d)
  1813.     else:
  1814.         return masked_array(d, mask=Numeric.ravel(m))
  1815.  
  1816. def concatenate (arrays, axis=0):
  1817.     "Concatenate the arrays along the given axis"
  1818.     d = []
  1819.     for x in arrays:
  1820.         d.append(filled(x))
  1821.     d = Numeric.concatenate(d, axis)
  1822.     for x in arrays:
  1823.         if getmask(x) is not None: break
  1824.     else:
  1825.         return masked_array(d)
  1826.     dm = []
  1827.     for x in arrays:
  1828.         dm.append(getmaskarray(x))
  1829.     dm = Numeric.concatenate(dm, axis)
  1830.     return masked_array(d, mask=dm)
  1831.  
  1832. def take (a, indices, axis=0):
  1833.     "take(a, indices, axis=0) returns selection of items from a."
  1834.     m = getmask(a)
  1835.     d = masked_array(a).raw_data()
  1836.     if m is None:
  1837.         return masked_array(Numeric.take(d, indices, axis))
  1838.     else:
  1839.         return masked_array(Numeric.take(d, indices, axis), 
  1840.                      mask = Numeric.take(m, indices, axis))
  1841.  
  1842. def transpose(a, axes=None):
  1843.     "transpose(a, axes=None) reorder dimensions per tuple axes"
  1844.     m = getmask(a)
  1845.     d = filled(a)
  1846.     if m is None:
  1847.         return masked_array(Numeric.transpose(d, axes))
  1848.     else:
  1849.         return masked_array(Numeric.transpose(d, axes), 
  1850.                      mask = Numeric.transpose(m, axes))
  1851.     
  1852. def put (a, indices, values):
  1853.     "put(a, indices, values) sets storage-indexed locations to corresponding values. values and indices are filled if necessary."
  1854.     d = a.raw_data()
  1855.     ind = filled(indices)
  1856.     v = filled(values)
  1857.     Numeric.put (d, ind, v)
  1858.     m = getmask(a)
  1859.     if m is not None:
  1860.         a.unshare_mask()
  1861.         Numeric.put(a.raw_mask(), ind, 0)
  1862.  
  1863. def putmask (a, mask, values):
  1864.     "put (a, mask, values) sets a where mask is true."
  1865.     if mask is None:
  1866.         return
  1867.     Numeric.putmask(a.raw_data(), mask, values)
  1868.     m = getmask(a)
  1869.     if m is None: return
  1870.     a.unshare_mask()
  1871.     Numeric.putmask(a.raw_mask(), mask, 0)
  1872.  
  1873. def innerproduct(a,b):
  1874.     """innerproduct(a,b) returns the dot product of two arrays, which has
  1875.     shape a.shape[:-1] + b.shape[:-1] with elements computed by summing the
  1876.     product of the elements from the last dimensions of a and b.
  1877.     Masked elements are replace by zeros.
  1878.     """
  1879.     fa = filled(a, 0)
  1880.     fb = filled(b, 0)
  1881.     if len(fa.shape) == 0: fa.shape = (1,)
  1882.     if len(fb.shape) == 0: fb.shape = (1,)
  1883.     return masked_array(Numeric.innerproduct(fa, fb))
  1884.  
  1885. def outerproduct(a, b):
  1886.     """outerproduct(a,b) = {a[i]*b[j]}, has shape (len(a),len(b))"""
  1887.     fa = filled(a,0).flat
  1888.     fb = filled(b,0).flat
  1889.     d = Numeric.outerproduct(fa, fb)
  1890.     ma = getmask(a)
  1891.     mb = getmask(b)
  1892.     if ma is None and mb is None:
  1893.         return masked_array(d)
  1894.     ma = getmaskarray(a)
  1895.     mb = getmaskarray(b)
  1896.     m = make_mask(1-Numeric.outerproduct(1-ma,1-mb), copy=0)
  1897.     return masked_array(d, m)
  1898.  
  1899. def dot(a, b):
  1900.     """dot(a,b) returns matrix-multiplication between a and b.  The product-sum
  1901.     is over the last dimension of a and the second-to-last dimension of b.
  1902.     Masked values are replaced by zeros. See also innerproduct.
  1903.     """
  1904.     return innerproduct(filled(a,0), Numeric.swapaxes(filled(b,0), -1, -2))
  1905.  
  1906. def compress(condition, x, dimension=-1):
  1907.     """Select those parts of x for which condition is true.
  1908.        Masked values in condition are considered false.
  1909.     """
  1910.     c = filled(condition, 0)
  1911.     m = getmask(x)
  1912.     if m is not None:
  1913.         m=Numeric.compress(c, m, dimension)
  1914.     d = Numeric.compress(c, filled(x), dimension)
  1915.     return masked_array(d, m)
  1916.  
  1917. class _minimum_operation:
  1918.     "Object to calculate minima"
  1919.     def __init__ (self):
  1920.         """minimum(a, b) or minimum(a)
  1921.            In one argument case returns the scalar minimum.
  1922.         """
  1923.         pass
  1924.  
  1925.     def __call__ (self, a, b=None):
  1926.         "Execute the call behavior."
  1927.         if b is None:
  1928.             m = getmask(a)
  1929.             if m is None: 
  1930.                 d = min(filled(a).flat)
  1931.                 return d
  1932.             ac = a.compressed()
  1933.             if len(ac) == 0:
  1934.                 return masked
  1935.             else:
  1936.                 return min(ac.raw_data())
  1937.         else:
  1938.             return where(less(a, b), a, b)[...]
  1939.        
  1940.     def reduce (self, target, axis=0):
  1941.         """Reduce target along the given axis."""
  1942.         m = getmask(target)
  1943.         if m is None:
  1944.             t = filled(target)
  1945.             return masked_array (Numeric.minimum.reduce (t, axis))
  1946.         else:
  1947.             t = Numeric.minimum.reduce(filled(target, minimum_fill_value(target)), axis)
  1948.             m = Numeric.logical_and.reduce(m, axis)
  1949.             return masked_array(t, m, fill_value(target))
  1950.  
  1951.     def outer (self, a, b):
  1952.         "Return the function applied to the outer product of a and b."
  1953.         ma = getmask(a)
  1954.         mb = getmask(b)
  1955.         if ma is None and mb is None:
  1956.             m = None
  1957.         else:
  1958.             ma = getmaskarray(a)
  1959.             mb = getmaskarray(b)
  1960.             m = logical_or.outer(ma, mb)
  1961.         d = Numeric.minimum.outer(filled(a), filled(b))
  1962.         return masked_array(d, m)
  1963.  
  1964. minimum = _minimum_operation () 
  1965.     
  1966. class _maximum_operation:
  1967.     "Object to calculate maxima"
  1968.     def __init__ (self):
  1969.         """maximum(a, b) or maximum(a)
  1970.            In one argument case returns the scalar maximum.
  1971.         """
  1972.         pass
  1973.  
  1974.     def __call__ (self, a, b=None):
  1975.         "Execute the call behavior."
  1976.         if b is None:
  1977.             m = getmask(a)
  1978.             if m is None: 
  1979.                 d = max(filled(a).flat)
  1980.                 return d
  1981.             ac = a.compressed()
  1982.             if len(ac) == 0:
  1983.                 return masked
  1984.             else:
  1985.                 return max(ac.raw_data())
  1986.         else:
  1987.             return where(greater(a, b), a, b)[...]
  1988.        
  1989.     def reduce (self, target, axis=0):
  1990.         """Reduce target along the given axis."""
  1991.         m = getmask(target)
  1992.         if m is None:
  1993.             t = filled(target)
  1994.             return masked_array (Numeric.maximum.reduce (t, axis))
  1995.         else:
  1996.             t = Numeric.maximum.reduce(filled(target, maximum_fill_value(target)), axis)
  1997.             m = Numeric.logical_and.reduce(m, axis)
  1998.             return masked_array(t, m, fill_value(target))
  1999.  
  2000.     def outer (self, a, b):
  2001.         "Return the function applied to the outer product of a and b."
  2002.         ma = getmask(a)
  2003.         mb = getmask(b)
  2004.         if ma is None and mb is None:
  2005.             m = None
  2006.         else:
  2007.             ma = getmaskarray(a)
  2008.             mb = getmaskarray(b)
  2009.             m = logical_or.outer(ma, mb)
  2010.         d = Numeric.maximum.outer(filled(a), filled(b))
  2011.         return masked_array(d, m)
  2012.  
  2013. maximum = _maximum_operation () 
  2014.     
  2015. def sort (x, axis = -1, fill_value=None):
  2016.     """If x does not have a mask, return a masked array formed from the 
  2017.        result of Numeric.sort(x, axis).
  2018.        Otherwise, fill x with fill_value. Sort it. 
  2019.        Set a mask where the result is equal to fill_value.
  2020.        Note that this may have unintended consequences if the data contains the
  2021.        fill value at a non-masked site.
  2022.  
  2023.        If fill_value is not given the default fill value for x's type will be
  2024.        used.
  2025.     """
  2026.     if fill_value is None:
  2027.         fill_value = default_fill_value (x)      
  2028.     d = filled(x, fill_value)
  2029.     s = Numeric.sort(d, axis)
  2030.     if getmask(x) is None:
  2031.         return masked_array(s)
  2032.     return masked_values(s, fill_value, copy=0)
  2033.  
  2034. def diagonal(a, k = 0, axis1=0, axis2=1):
  2035.     """diagonal(a,k=0,axis1=0, axis2=1) = the k'th diagonal of a"""
  2036.     d = Numeric.diagonal(filled(a), k, axis1, axis2)
  2037.     m = getmask(a)
  2038.     if m is None:
  2039.         return masked_array(d, m)
  2040.     else:
  2041.         return masked_array(d, Numeric.diagonal(m, k, axis1, axis2))
  2042.     
  2043. def argsort (x, axis = -1, fill_value=None):
  2044.     """Treating masked values as if they have the value fill_value,
  2045.        return sort indices for sorting along given axis.
  2046.        if fill_value is None, use fill_value(x)
  2047.        Returns a Numeric array.
  2048.     """        
  2049.     d = filled(x, fill_value)
  2050.     return Numeric.argsort(d, axis)
  2051.  
  2052. def argmin (x, axis = -1, fill_value=None):
  2053.     """Treating masked values as if they have the value fill_value,
  2054.        return indices for minimum values along given axis.
  2055.        if fill_value is None, use fill_value(x).
  2056.        Returns a Numeric array if x has more than one dimension.
  2057.        Otherwise, returns a scalar index.
  2058.     """        
  2059.     d = filled(x, fill_value)
  2060.     return Numeric.argmin(d, axis)
  2061.  
  2062. def argmax (x, axis = -1, fill_value=None):
  2063.     """Treating masked values as if they have the value fill_value,
  2064.        return sort indices for maximum along given axis.
  2065.        if fill_value is None, use -fill_value(x) if it exists.
  2066.        Returns a Numeric array if x has more than one dimension.
  2067.        Otherwise, returns a scalar index.
  2068.     """   
  2069.     if fill_value is None:
  2070.         fill_value = default_fill_value (x) 
  2071.         try:
  2072.             fill_value = - fill_value
  2073.         except:
  2074.             pass    
  2075.     d = filled(x, fill_value)
  2076.     return Numeric.argmax(d, axis)
  2077.  
  2078. def fromfunction (f, s):
  2079.     """apply f to s to create array as in Numeric."""
  2080.     return masked_array(Numeric.fromfunction(f,s)) 
  2081.  
  2082. def asarray(data, typecode=None):
  2083.     """asarray(data, typecode=None) = array(data, typecode=None, copy=0)
  2084.        Returns data if typecode if data is a MaskedArray and typecode None
  2085.        or the same.
  2086.     """
  2087.     if isinstance(data, MaskedArray) and \
  2088.         (typecode is None or typecode == data.typecode()):
  2089.         return data
  2090.     return array(data, typecode=typecode, copy=0)
  2091.     
  2092. # This section is stolen from a post about how to limit array printing.
  2093. __MaxElements = 300     #Maximum size for printing
  2094.  
  2095. def limitedArrayRepr(a, max_line_width = None, precision = None, suppress_small = None):
  2096.     "Calculate string representation, limiting size of output."
  2097.     global __MaxElements
  2098.     s = a.shape
  2099.     elems =  Numeric.multiply.reduce(s)
  2100.     if elems > __MaxElements:
  2101.         if len(s) > 1:
  2102.             return 'array (%s) , type = %s, has %d elements' % \
  2103.                  (string.join(map(str, s), ","), a.typecode(), elems)
  2104.         else:
  2105.             return Numeric.array2string (a[:__MaxElements], max_line_width, precision,
  2106.                  suppress_small,',',0) + \
  2107.                ('\n + %d more elements' % (elems - __MaxElements))
  2108.     else:
  2109.         return Numeric.array2string (a, max_line_width, precision, 
  2110.                 suppress_small,',',0)
  2111.  
  2112. __original_str = Numeric.array_str
  2113. __original_repr = Numeric.array_repr
  2114.  
  2115. def set_print_limit (m=0):
  2116.     "Set the maximum # of elements for printing arrays. <=0  = no limit"
  2117.     import multiarray
  2118.     global __MaxElements
  2119.     n = int(m)
  2120.     __MaxElements = n
  2121.     if n <= 0:
  2122.         Numeric.array_str = __original_str
  2123.         Numeric.array_repr = __original_repr
  2124.         multiarray.set_string_function(__original_str, 0)
  2125.         multiarray.set_string_function(__original_repr, 1)
  2126.     else:
  2127.         Numeric.array_str = limitedArrayRepr
  2128.         Numeric.array_repr = limitedArrayRepr
  2129.         multiarray.set_string_function(limitedArrayRepr, 0)
  2130.         multiarray.set_string_function(limitedArrayRepr, 1)
  2131.  
  2132. def get_print_limit ():
  2133.     "Get the maximum # of elements for printing arrays. "
  2134.     return __MaxElements
  2135.  
  2136. set_print_limit(__MaxElements)
  2137.  
  2138.